#!/usr/bin/env python3

import os
import sys
import glob
import re
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.patches import Rectangle
from pathlib import Path
from typing import Dict, List, Any, Tuple

# --- 1. 配置路径 ---
BASE_DIR = Path("/Users/chenxi/Documents/比赛/CCF/checker/example")
IN_DIR = BASE_DIR / "infile"
OUT_DIR = BASE_DIR / "outfile"
FIG_DIR = BASE_DIR / "fig"

# --- 2. 辅助函数：解析文件 ---

def parse_infile(filepath: Path) -> Dict[str, Any]:
    """解析 infile，返回 L, M, N 和一个请求字典。"""
    requests = {}
    with open(filepath, 'r') as f:
        line = f.readline().split()
        L, M, N = int(line[0]), int(line[1]), int(line[2])
        
        for i in range(N):
            line = f.readline().split()
            requests[i] = {
                'addr': int(line[0]),
                'size': int(line[1]),
                'start': int(line[2]),
                'time': int(line[3]),
                'id': i
            }
    return {'L': L, 'M': M, 'N': N, 'requests': requests}

def parse_outfile(filepath: Path) -> List[Dict[str, Any]]:
    """解析 outfile，返回一个操作列表。"""
    ops = []
    with open(filepath, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            
            op_name = parts[0]
            T = int(parts[1])
            op = {'op': op_name, 'T': T}
            
            if op_name in ("Reload", "Offload"):
                op['A'] = int(parts[2])
                op['S'] = int(parts[3])
            elif op_name == "Visit":
                op['A'] = int(parts[2]) # 这是任务 ID
            
            ops.append(op)
    return ops


# --- 3. [NEW] 变长时间轴映射函数 ---

def create_variable_timescale(events: List[Tuple[int, int]], 
                               compression_threshold: int = 1000,
                               min_segment_display: int = 50) -> Tuple[callable, callable, List[Tuple[int, int]]]:
    """
    创建变长时间轴映射。
    
    参数:
        events: [(start_time, end_time), ...] 所有事件的时间区间
        compression_threshold: 超过这个长度的空闲段会被压缩
        min_segment_display: 压缩后空闲段的最小显示宽度
    
    返回:
        (real_to_display, display_to_real, segments)
        - real_to_display: 真实时间 -> 显示时间的映射函数
        - display_to_real: 显示时间 -> 真实时间的映射函数  
        - segments: [(real_start, real_end, display_start, display_end, is_compressed), ...]
    """
    if not events:
        # 空事件列表，返回恒等映射
        identity = lambda x: x
        return identity, identity, []
    
    # 1. 收集所有时间点并排序
    time_points = set()
    for start, end in events:
        time_points.add(start)
        time_points.add(end)
    
    sorted_points = sorted(time_points)
    if not sorted_points:
        identity = lambda x: x
        return identity, identity, []
    
    # 2. 识别活跃段和空闲段
    segments = []
    display_cursor = 0
    
    for i in range(len(sorted_points)):
        real_start = sorted_points[i]
        real_end = sorted_points[i+1] if i+1 < len(sorted_points) else sorted_points[i]
        
        if real_start == real_end:
            continue
        
        # 检查这个时间段是否有活动
        has_activity = any(start <= real_start < end for start, end in events)
        
        segment_length = real_end - real_start
        
        if has_activity or segment_length <= compression_threshold:
            # 活跃段或短段：正常显示
            display_length = segment_length
            is_compressed = False
        else:
            # 长空闲段：压缩
            display_length = min_segment_display
            is_compressed = True
        
        display_end = display_cursor + display_length
        segments.append((real_start, real_end, display_cursor, display_end, is_compressed))
        display_cursor = display_end
    
    # 3. 创建映射函数
    def real_to_display(t_real: float) -> float:
        """真实时间 -> 显示时间"""
        for r_start, r_end, d_start, d_end, _ in segments:
            if r_start <= t_real <= r_end:
                # 线性插值
                ratio = (t_real - r_start) / (r_end - r_start) if r_end > r_start else 0
                return d_start + ratio * (d_end - d_start)
        # 超出范围，返回最后的显示时间
        if segments:
            return segments[-1][3]
        return t_real
    
    def display_to_real(t_display: float) -> float:
        """显示时间 -> 真实时间"""
        for r_start, r_end, d_start, d_end, _ in segments:
            if d_start <= t_display <= d_end:
                ratio = (t_display - d_start) / (d_end - d_start) if d_end > d_start else 0
                return r_start + ratio * (r_end - r_start)
        if segments:
            return segments[-1][1]
        return t_display
    
    return real_to_display, display_to_real, segments


# --- 4. [MODIFIED] 甘特图可视化函数 ---

def visualize_case(infile_path: Path, outfile_path: Path, fig_path: Path):
    """
    为单个用例生成并保存甘特图可视化。
    [!!! MODIFIED !!!] 使用三个独立的变长时间轴。
    """
    
    try:
        in_data = parse_infile(infile_path)
        out_data = parse_outfile(outfile_path)
    except Exception as e:
        print(f"  [!] 失败: 解析文件时出错 {e}")
        return

    L, M, N = in_data['L'], in_data['M'], in_data['N']
    requests = in_data['requests']
    
    # 模拟 HBM 状态
    try:
        hbm_state = np.zeros(L, dtype=np.int8)
    except MemoryError:
        print(f"  [!] 失败: 无法为 L={L} 分配内存。跳过...")
        return
    
    # 收集三类事件
    request_events = []  # (task_id, real_start, real_duration, label, color)
    visit_events = []
    io_events = []
    
    colors = list(mcolors.TABLEAU_COLORS.values())
    
    # 收集 Requests (理论时间)
    for task_id, req in requests.items():
        color = colors[task_id % len(colors)]
        request_events.append((
            task_id,
            req['start'],
            req['time'],
            f"R{task_id}",
            color
        ))
    
    # 遍历 Outfile 收集 Visits 和 IO
    T_fin = 0
    for op in out_data:
        T_start = op['T']
        T_fin = max(T_fin, T_start)
        
        if op['op'] == 'Reload':
            A, S = op['A'], op['S']
            try:
                chunk = hbm_state[A : A+S]
                new_bytes = (chunk == 0)
                actual_reload_size = int(np.sum(new_bytes))
                chunk[new_bytes] = 1 
            except ValueError:
                continue 

            duration = 40 * actual_reload_size
            if duration > 0:
                io_events.append((T_start, duration, 'R', colors[0]))
                T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Offload':
            A, S = op['A'], op['S']
            try:
                chunk = hbm_state[A : A+S]
                offload_bytes = (chunk == 1)
                actual_offload_size = int(np.sum(offload_bytes))
                chunk[offload_bytes] = 0
            except ValueError:
                continue

            duration = 40 * actual_offload_size
            if duration > 0:
                io_events.append((T_start, duration, 'O', colors[3]))
                T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Visit':
            task_id = op['A']
            if task_id not in requests:
                continue
                
            req = requests[task_id]
            duration = req['time']
            color = colors[task_id % len(colors)]
            
            visit_events.append((task_id, T_start, duration, f"V{task_id}", color))
            T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Fin':
            T_fin = max(T_fin, T_start)
            break
    
    # --- 5. 创建三个独立的变长时间轴 ---
    
    # Requests 时间轴
    request_intervals = [(start, start + dur) for _, start, dur, _, _ in request_events]
    req_r2d, req_d2r, req_segments = create_variable_timescale(request_intervals)
    
    # Visits 时间轴
    visit_intervals = [(start, start + dur) for _, start, dur, _, _ in visit_events]
    vis_r2d, vis_d2r, vis_segments = create_variable_timescale(visit_intervals)
    
    # IO 时间轴
    io_intervals = [(start, start + dur) for start, dur, _, _ in io_events]
    io_r2d, io_d2r, io_segments = create_variable_timescale(io_intervals)
    
    # --- 6. 转换事件到显示坐标 ---
    
    request_bars = []
    for task_id, real_start, real_dur, label, color in request_events:
        disp_start = req_r2d(real_start)
        disp_end = req_r2d(real_start + real_dur)
        disp_dur = disp_end - disp_start
        request_bars.append((task_id, disp_start, disp_dur, label, color, real_start, real_dur))
    
    visit_bars = []
    for task_id, real_start, real_dur, label, color in visit_events:
        disp_start = vis_r2d(real_start)
        disp_end = vis_r2d(real_start + real_dur)
        disp_dur = disp_end - disp_start
        visit_bars.append((task_id, disp_start, disp_dur, label, color, real_start, real_dur))
    
    io_bars = []
    for real_start, real_dur, label, color in io_events:
        disp_start = io_r2d(real_start)
        disp_end = io_r2d(real_start + real_dur)
        disp_dur = disp_end - disp_start
        io_bars.append((disp_start, disp_dur, label, color, real_start, real_dur))
    
    # --- 7. 绘图 ---
    
    height_ratio = [max(1, N), max(1, N), 1.5]
    
    fig, (ax_requests, ax_visits, ax_io) = plt.subplots(
        3, 1,
        figsize=(24, 3 + N * 1.0 + 1.0),
        gridspec_kw={'height_ratios': height_ratio}
    )
    
    # --- 子图 1: Requests ---
    ax_requests.set_ylabel("Requests (Input)", fontsize=12, weight='bold')
    
    for (task_id, disp_start, disp_dur, label, color, real_start, real_dur) in request_bars:
        ax_requests.barh(
            y=task_id, 
            width=disp_dur, 
            left=disp_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.4,
            color=color,
            hatch='//'
        )
        # 标签显示真实时间
        ax_requests.text(
            disp_start + disp_dur/2, 
            task_id, 
            f"{label}\n[{real_start}-{real_start+real_dur}]", 
            ha='center', 
            va='center', 
            color='black', 
            alpha=0.8,
            fontsize=8
        )
    
    ax_requests.set_yticks(range(N))
    ax_requests.set_yticklabels([f"R{i}" for i in range(N)])
    ax_requests.set_ylim(N - 0.5, -0.5)
    ax_requests.grid(axis='x', linestyle=':', alpha=0.6)
    ax_requests.set_title("Requests Timeline (Variable Scale)", fontsize=10, style='italic')
    
    # 添加压缩段标记
    for r_start, r_end, d_start, d_end, is_compressed in req_segments:
        if is_compressed:
            ax_requests.axvspan(d_start, d_end, alpha=0.1, color='gray', zorder=0)
    
    # --- 子图 2: Visits ---
    ax_visits.set_ylabel("Visits (Scheduled)", fontsize=12, weight='bold')
    
    for (task_id, disp_start, disp_dur, label, color, real_start, real_dur) in visit_bars:
        ax_visits.barh(
            y=task_id, 
            width=disp_dur, 
            left=disp_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.8,
            color=color
        )
        ax_visits.text(
            disp_start + disp_dur/2, 
            task_id, 
            f"{label}\n[{real_start}-{real_start+real_dur}]", 
            ha='center', 
            va='center', 
            color='white', 
            weight='bold',
            fontsize=8
        )
    
    ax_visits.set_yticks(range(N))
    ax_visits.set_yticklabels([f"V{i}" for i in range(N)])
    ax_visits.set_ylim(N - 0.5, -0.5)
    ax_visits.grid(axis='x', linestyle=':', alpha=0.6)
    ax_visits.set_title("Visits Timeline (Variable Scale)", fontsize=10, style='italic')
    
    # 添加压缩段标记
    for r_start, r_end, d_start, d_end, is_compressed in vis_segments:
        if is_compressed:
            ax_visits.axvspan(d_start, d_end, alpha=0.1, color='gray', zorder=0)
    
    # --- 子图 3: IO ---
    ax_io.set_ylabel("IO (Serial)", fontsize=12, weight='bold')
    
    for (disp_start, disp_dur, label, color, real_start, real_dur) in io_bars:
        ax_io.barh(
            y=0, 
            width=disp_dur, 
            left=disp_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.75,
            color=color
        )
        # 标签
        ax_io.text(
            disp_start + disp_dur/2, 
            0, 
            f"{label}\n[{real_start}-{real_start+real_dur}]", 
            ha='center', 
            va='center', 
            color='white', 
            weight='bold',
            fontsize=8
        )
    
    ax_io.set_yticks([0])
    ax_io.set_yticklabels(["IO Lane"])
    ax_io.set_ylim(-0.5, 0.5)
    ax_io.grid(axis='x', linestyle=':', alpha=0.6)
    ax_io.set_xlabel("Display Time (Variable Scale)", fontsize=12, weight='bold')
    ax_io.set_title("IO Timeline (Variable Scale)", fontsize=10, style='italic')
    
    # 添加压缩段标记
    for r_start, r_end, d_start, d_end, is_compressed in io_segments:
        if is_compressed:
            ax_io.axvspan(d_start, d_end, alpha=0.1, color='gray', zorder=0)
    
    # 设置 X 轴范围（每个子图独立）
    if request_bars:
        max_disp_req = max(disp_start + disp_dur for _, disp_start, disp_dur, _, _, _, _ in request_bars)
        ax_requests.set_xlim(left=0, right=max_disp_req * 1.02)
    
    if visit_bars:
        max_disp_vis = max(disp_start + disp_dur for _, disp_start, disp_dur, _, _, _, _ in visit_bars)
        ax_visits.set_xlim(left=0, right=max_disp_vis * 1.02)
    
    if io_bars:
        max_disp_io = max(disp_start + disp_dur for disp_start, disp_dur, _, _, _, _ in io_bars)
        ax_io.set_xlim(left=0, right=max_disp_io * 1.02)
    
    # --- 8. 保存 ---
    plt.tight_layout(rect=[0, 0.03, 1, 1.0])
    plt.savefig(fig_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    print(f"  [✓] 成功: 已保存到 {fig_path.name}")


# --- 9. 主程序 ---

def main():
    print(f"--- 开始可视化 (变长时间轴甘特图) ---")
    print(f"输入目录: {IN_DIR}")
    print(f"输出目录: {OUT_DIR}")
    print(f"图形目录: {FIG_DIR}")
    print("--------------------")

    FIG_DIR.mkdir(parents=True, exist_ok=True)
    
    infile_paths = sorted(glob.glob(str(IN_DIR / "infile_*.txt")))
    
    if not infile_paths:
        print(f"[!] 错误: 在 {IN_DIR} 中未找到 'infile_*.txt' 文件。")
        return

    for infile_path in infile_paths:
        infile_path = Path(infile_path)
        
        match = re.search(r'infile_(\w+)\.txt', infile_path.name)
        if not match:
            print(f"[!] 跳过: 无法解析文件名 {infile_path.name}")
            continue
            
        file_id = match.group(1)
        outfile_name = f"outfile_{file_id}.txt"
        outfile_path = OUT_DIR / outfile_name
        
        fig_name = f"fig_{file_id}_gantt_variable.png"
        fig_path = FIG_DIR / fig_name
        
        print(f"[*] 正在处理: {infile_path.name} -> {outfile_name}")

        if not outfile_path.exists():
            print(f"  [!] 跳过: 未找到对应的 {outfile_name}")
            continue
            
        try:
            visualize_case(infile_path, outfile_path, fig_path)
        except Exception as e:
            print(f"  [!] 失败: 处理 {infile_path.name} 时发生意外错误: {e}")
            import traceback
            traceback.print_exc()

    print("--------------------")
    print("--- 可视化完成 ---")

if __name__ == "__main__":
    main()
